import sys
from mat import write_tsv
from tqdm import tqdm
import pdb
IN_PATH = sys.argv[1]
OUT_PATH = sys.argv[2]

keep_tags = ['n', 'v', 'a', 'd', 'f'] # 04:18
# keep_tags = ['n', 'v', 'd', 'a', 'f', 'ns', 'vn'] # 03:18

with open(IN_PATH, encoding='gbk') as f:
    corpus = f.readlines()


def split_and_filter(text):
    if len(text) < 24:
        return []

    def word_tag(raw_seg):
        seg = raw_seg.split('/')
        assert len(seg) == 2
        return seg

    def split(text):
        text = text[23:].strip().split()
        text = [word_tag(seg) for seg in text]
        return text

    return [word for word, tag in split(text) for t in keep_tags if tag.startswith(t)]
    # return [word for word, tag in split(text) if tag in keep_tags]

docs = []
for line in corpus:
    words = split_and_filter(line)
    if len(words):
        docs.append(words)
print(f'num of docs: {len(docs)}')

vocab = set()
for doc in docs:
    for word in doc:
        vocab.add(word)
print(f'size of vocab: {len(vocab)}')

def bow():
    for doc in tqdm(docs):
        cnt = {}
        for word in doc:
            cnt[word] = cnt.get(word, 0) + 1
        yield [cnt.get(w, 0) for w in vocab]

write_tsv(bow(), OUT_PATH)
